/**
 * Copyright 2019-2022 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#ifndef FRAMEWORK_TENSOR_AIPP_PARA_IMPL_H
#define FRAMEWORK_TENSOR_AIPP_PARA_IMPL_H
#include "tensor/aipp_para.h"
#include "framework/c/hiai_tensor_aipp_para.h"
#include "tensor/repo/compatible/hiai_tensor_aipp_para_def.h"

namespace hiai {
using TensorAippCommPara = std::pair<HIAI_MR_TensorAippCommPara*, uint8_t>;

class AIPPParaBufferImpl {
public:
    static HIAI_MR_TensorAippPara* Init(uint32_t batchCount);
    static Status Release(HIAI_MR_TensorAippPara* paraBuff);

    static void* GetRawBuffer(HIAI_MR_TensorAippPara* paraBuff);
    static uint32_t GetRawBufferSize(HIAI_MR_TensorAippPara* paraBuff);

    static int GetInputIndex(HIAI_MR_TensorAippPara* paraBuff);
    static void SetInputIndex(HIAI_MR_TensorAippPara* paraBuff, uint32_t inputIndex);
    static int GetInputAippIndex(HIAI_MR_TensorAippPara* paraBuff);
    static void SetInputAippIndex(HIAI_MR_TensorAippPara* paraBuff, uint32_t inputAippIndex);

    static Status SetCscPara(HIAI_MR_TensorAippPara* paraBuff, ImageFormat inputFormat,
        ImageFormat targetFormat, ImageColorSpace colorSpace);
    static Status SetCscMatrixPara(HIAI_MR_TensorAippPara* paraBuff, CscMatrixPara cscMatrixPara);
    static Status GetCscPara(HIAI_MR_TensorAippPara* paraBuff, ImageFormat* inputFormat,
        ImageFormat* targetFormat, ImageColorSpace* colorSpace);
    static Status GetCscMatrixPara(HIAI_MR_TensorAippPara* paraBuff, CscMatrixPara* cscMatrixPara);

    static Status SetChannelSwapPara(HIAI_MR_TensorAippPara* paraBuff, bool rbuvSwapSwitch, bool axSwapSwitch);
    static Status GetChannelSwapPara(HIAI_MR_TensorAippPara* paraBuff, bool* rbuvSwapSwitch, bool* axSwapSwitch);

    static Status SetSingleBatchMultiCrop(HIAI_MR_TensorAippPara* paraBuff, bool singleBatchMutiCrop);
    static bool GetSingleBatchMultiCrop(HIAI_MR_TensorAippPara* paraBuff);

    static uint32_t GetBatchCount(HIAI_MR_TensorAippPara* paraBuff);

    static Status SetInputFormat(HIAI_MR_TensorAippPara* tensorAippPara, ImageFormat inputFormat);
    static ImageFormat GetInputFormat(HIAI_MR_TensorAippPara* tensorAippPara);

    static Status GetInputShape(HIAI_MR_TensorAippPara* tensorAippPara,
        uint32_t* srcImageW, uint32_t* srcImageH);
    static Status SetInputShape(HIAI_MR_TensorAippPara* tensorAippPara,
        uint32_t srcImageW, uint32_t srcImageH);

    static Status SetCropPos(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex,
        uint32_t cropStartPosW, uint32_t cropStartPosH);
    static Status SetCropSize(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex,
        uint32_t cropSizeW, uint32_t cropSizeH);
    static Status GetCropPos(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex,
    uint32_t* cropStartPosW, uint32_t* cropStartPosH);
    static Status GetCropSize(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex,
        uint32_t* cropSizeW, uint32_t* cropSizeH);

    static Status SetResizePara(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex,
        uint32_t resizeOutputSizeW, uint32_t resizeOutputSizeH);
    static Status GetResizePara(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex,
        uint32_t* resizeOutputSizeW, uint32_t* resizeOutputSizeH);

    static Status SetPadTopSize(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex,
    uint32_t paddingSizeTop);
    static Status SetPadBottomSize(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex,
        uint32_t paddingSizeBottom);
    static Status SetPadLeftSize(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex,
        uint32_t paddingSizeLeft);
    static Status SetPadRightSize(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex,
        uint32_t paddingSizeRight);
    static Status SetPadChannelValue(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex,
        uint32_t chnValue[], uint32_t chnNum);
    static Status GetPadTopSize(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex,
        uint32_t* paddingSizeTop);
    static Status GetPadBottomSize(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex,
        uint32_t* paddingSizeBottom);
    static Status GetPadLeftSize(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex,
        uint32_t* paddingSizeLeft);
    static Status GetPadRightSize(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex,
        uint32_t* paddingSizeRight);
    static Status GetPadChannelValue(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex,
        uint32_t chnValue[], uint32_t chnNum);

    static Status SetDtcPixelMeanPara(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex,
        int32_t pixelMeanPara[], uint32_t chnNum);
    static Status SetDtcPixelMinPara(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex,
        float pixelMinPara[], uint32_t chnNum);
    static Status SetDtcPixelVarReciPara(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex,
        float pixelVarReciPara[], uint32_t chnNum);
    static Status GetDtcPixelMeanPara(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex,
        int32_t pixelMeanPara[], uint32_t chnNum);
    static Status GetDtcPixelMinPara(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex,
        float pixelMinPara[], uint32_t chnNum);
    static Status GetDtcPixelVarReciPara(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex,
        float pixelVarReciPara[], uint32_t chnNum);

    static Status SetRotatePara(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex, float rotateAngle);
    static Status GetRotatePara(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex, float* rotateAngle);

    static bool GetCropSwitch(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex);
    static bool GetResizeSwitch(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex);
    static bool GetCscSwitch(HIAI_MR_TensorAippPara* tensorAippPara);
    static bool GetPaddingSwitch(HIAI_MR_TensorAippPara* tensorAippPara, uint32_t batchIndex);
};

class AIPPParaImpl : public IAIPPPara {
public:
    AIPPParaImpl() = default;
    ~AIPPParaImpl() override;
    AIPPParaImpl(const AIPPParaImpl&) = delete;
    AIPPParaImpl& operator=(const AIPPParaImpl&) = delete;

    Status Init(uint32_t batchCount);
    Status Init(HIAI_MR_TensorAippPara* paraBuff);

    uint32_t GetBatchCount() override;

    Status SetInputIndex(uint32_t inputIndex) override;
    int32_t GetInputIndex() override;
    Status SetInputAippIndex(uint32_t inputAippIndex) override;
    int32_t GetInputAippIndex() override;

    Status SetCscPara(ImageFormat targetFormat, ImageColorSpace colorSpace) override;
    Status SetCscPara(CscMatrixPara cscPara) override;
    CscMatrixPara GetCscMatrixPara() override;
    CscPara GetCscPara() override;

    Status SetChannelSwapPara(ChannelSwapPara&& channelSwapPara) override;
    ChannelSwapPara GetChannelSwapPara() override;

    Status SetSingleBatchMultiCrop(bool singleBatchMutiCrop) override;
    bool GetSingleBatchMultiCrop() override;
    // end common interface

    Status SetCropPara(CropPara&& cropPara) override;
    Status SetCropPara(uint32_t batchIndex, CropPara&& cropPara) override;
    CropPara GetCropPara(uint32_t batchIndex) override;

    Status SetResizePara(ResizePara&& resizePara) override;
    Status SetResizePara(uint32_t batchIndex, ResizePara&& resizePara) override;
    ResizePara GetResizePara(uint32_t batchIndex) override;

    Status SetPaddingPara(PadPara&& paddingPara) override;
    Status SetPaddingPara(uint32_t batchIndex, PadPara&& paddingPara) override;
    PadPara GetPaddingPara(uint32_t batchIndex) override;

    Status SetDtcPara(DtcPara&& dtcPara) override;
    Status SetDtcPara(uint32_t batchIndex, DtcPara&& dtcPara) override;
    DtcPara GetDtcPara(uint32_t batchIndex) override;

    Status SetRotatePara(RotatePara&& rotatePara) override;
    Status SetRotatePara(uint32_t batchIndex, RotatePara&& rotatePara) override;
    RotatePara GetRotatePara(uint32_t batchIndex) override;

    // internel
    Status SetInputFormat(ImageFormat inputFormat) override;
    ImageFormat GetInputFormat() override;
    std::vector<int32_t> GetInputShape() override;
    Status SetInputShape(std::vector<int32_t>& shape) override;

    void* GetData() override;

    size_t GetSize() const override;

    HIAI_TENSOR_API_EXPORT HIAI_MR_TensorAippPara* GetParaBuffer();

public:
    bool GetEnableCrop(uint32_t batchIndex);

    bool GetEnableResize(uint32_t batchIndex);

    bool GetEnableCsc();

    bool GetEnablePadding(uint32_t batchIndex);

protected:
    Status GetAippParaBufferImpl(std::shared_ptr<AIPPParaBufferImpl>& aippParaImpl);

private:
    Status InitAippPara(uint32_t batchCount);

    template <typename T, typename UpdateParaFunc>
    Status SetAippFuncPara(T&& funcPara, UpdateParaFunc updateParaFunc);

    template <typename T, typename UpdateParaFunc>
    Status SetAippFuncPara(uint32_t batchIndex, T&& funcPara, UpdateParaFunc updateParaFunc);

private:
    HIAI_MR_TensorAippPara* paraBuff_ {nullptr};
    void* rawBuffer_ {nullptr};
    ImageFormat targetFormat_ {ImageFormat::INVALID};
    ImageColorSpace colorSpace_ {ImageColorSpace::INVALID};
};
HIAI_TENSOR_API_EXPORT HIAI_MR_TensorAippPara* GetTensorAippParaFromAippPara(
    const std::shared_ptr<IAIPPPara>& aippPara);
} // namespace hiai
#endif